/* -*- P4_16 -*- */
#include <core.p4>
#include <v1model.p4>

// --- Protocol Constants ---
const bit<16> ETH_TYPE_IPV4 = 0x0800;
const bit<8>  IP_PROTO_TCP = 6;

/*************************************************************************
*** T Y P E D E F S   A N D   F R A M E S ***
*************************************************************************/

// --- Custom Type Aliases ---
typedef bit<48> DatalinkAddr;
typedef bit<32> InternetAddr;
typedef bit<16> TransportPort;
typedef bit<9>  EgressPortID;
typedef bit<14> FlowHashValue;

// --- Header Definitions (New PascalCase Naming & Reordered Fields) ---

header LinkFrame {
    DatalinkAddr destination;
    DatalinkAddr source;
    bit<16>      protocolId;
}

header NetworkHeader {
    bit<4>       ver;
    bit<4>       headerLength;
    bit<8>       tos;
    bit<16>      identification;
    bit<16>      totalLength;
    bit<3>       flags;
    bit<13>      fragOffset;
    bit<8>       protocol;
    bit<8>       timeToLive;
    bit<16>      hdrChecksum;
    InternetAddr srcAddress;
    InternetAddr dstAddress;
}

header TcpSegment {
    TransportPort sourcePort;
    TransportPort destPort;
    bit<32>       sequenceNum;
    bit<32>       ackNum;
    bit<4>        dataOffset;
    bit<16>       windowSize;
    bit<3>        res;
    bit<3>        ecn;
    bit<6>        controlBits;
    bit<16>       checksum;
    bit<16>       urgentPointer;
}

// --- Metadata and Header Structures ---

struct ProcessingMetadata {
    FlowHashValue hash_selector; // ECMP Hash result storage
}

struct PacketLayout {
    LinkFrame     link;
    NetworkHeader network;
    TcpSegment    transport;
}

/*************************************************************************
*** P A R S E R   U N I T ***
*************************************************************************/

parser FrameParser(packet_in p,
                    out PacketLayout layout,
                    inout ProcessingMetadata p_meta,
                    inout standard_metadata_t std_meta) {

    state Parse_Start {
        transition Parse_L2;
    }
    
    // New state naming and logic ordering
    state Parse_L2 {
        p.extract(layout.link);
        transition select(layout.link.protocolId) {
            ETH_TYPE_IPV4: Check_L3;
            default: accept;
        }
    }

    state Check_L3 {
        p.extract(layout.network);
        transition select(layout.network.protocol) {
            IP_PROTO_TCP: Parse_L4_TCP;
            default: accept;
        }
    }

    state Parse_L4_TCP {
        p.extract(layout.transport);
        transition accept;
    }
}

/*************************************************************************
*** V E R I F Y C H E C K S U M ***
*************************************************************************/

control VerifierBlock(inout PacketLayout layout, inout ProcessingMetadata p_meta) {
    apply { /* Verification logic placeholder */ }
}

/*************************************************************************
*** I N G R E S S   P R O C E S S O R ***
*************************************************************************/

control IngressProcessor(inout PacketLayout layout,
                         inout ProcessingMetadata p_meta,
                         inout standard_metadata_t std_meta) {

    // --- Actions ---
    action DropPacket() {
        mark_to_drop(std_meta);
    }

    // A1: Compute Hash for ECMP next-hop selection
    action CalculateHash(bit<16> hash_seed, bit<32> group_size) {
        hash(p_meta.hash_selector,
             HashAlgorithm.crc16,
             hash_seed,
             { layout.network.srcAddress,
               layout.network.dstAddress,
               layout.network.protocol,
               layout.transport.sourcePort,
               layout.transport.destPort },
             group_size);
    }

    // A2: Set ECMP next-hop and rewrite DMAC/DIP
    action SetEcmpRoute(DatalinkAddr next_hop_mac, InternetAddr next_hop_ip, EgressPortID target_port) {
        layout.link.destination = next_hop_mac;
        layout.network.dstAddress = next_hop_ip;
        std_meta.egress_spec = target_port;
        layout.network.timeToLive = layout.network.timeToLive - 1;
    }

    // A3: Set standard IPv4 forward route (no DIP change)
    action SetStandardRoute(DatalinkAddr next_hop_mac, EgressPortID target_port) {
        layout.link.destination = next_hop_mac;
        std_meta.egress_spec = target_port;
        layout.network.timeToLive = layout.network.timeToLive - 1;
    }

    // --- Tables ---

    // T1: ECMP Group Lookup (LPM)
    table RouteGroupTable {
        key = {
            layout.network.dstAddress: lpm;
        }
        actions = {
            CalculateHash;
            DropPacket;
        }
        size = 1024;
        default_action = DropPacket();
    }

    // T2: ECMP Next-Hop Selection (Exact Match on Hash)
    table EcmpNextHopTable {
        key = {
            p_meta.hash_selector: exact;
        }
        actions = {
            SetEcmpRoute;
            DropPacket;
        }
        size = 16;
        default_action = DropPacket();
    }

    // T3: Standard Route Lookup (LPM)
    table StandardRouteTable {
        key = {
            layout.network.dstAddress: lpm;
        }
        actions = {
            SetStandardRoute;
            DropPacket;
        }
        size = 1024;
        default_action = DropPacket();
    }

    // --- Apply Logic ---
    apply {
        if (layout.network.isValid() && layout.network.timeToLive > 0) {
            
            // Apply the first table and check if it resulted in a hit
            // Reversing the check order slightly: use action_run result directly.
            table_ecmp_group_result_t group_result = RouteGroupTable.apply();

            if (group_result.hit) {
                // If a hash was calculated, proceed to the next-hop selection
                EcmpNextHopTable.apply();
            } else {
                // Use the standard forwarding path
                StandardRouteTable.apply();
            }
        } else {
            DropPacket();
        }
    }
}

/*************************************************************************
*** E G R E S S   P R O C E S S O R ***
*************************************************************************/

control EgressProcessor(inout PacketLayout layout,
                        inout ProcessingMetadata p_meta,
                        inout standard_metadata_t std_meta) {

    action RewriteSMAC(DatalinkAddr device_smac) {
        layout.link.source = device_smac;
    }

    action DropPacket() {
        mark_to_drop(std_meta);
    }

    table EgressMacTable {
        key = {
            std_meta.egress_port: exact;
        }
        actions = {
            RewriteSMAC;
            DropPacket;
        }
        size = 256;
    }

    apply {
        EgressMacTable.apply();
    }
}

/*************************************************************************
*** C H E C K S U M   C O M P U T E R ***
*************************************************************************/

control ChecksumUpdater(inout PacketLayout layout, inout ProcessingMetadata p_meta) {
    apply {
        // Compute IPv4 header checksum based on modified header
        update_checksum(
            layout.network.isValid(),
            {
                layout.network.ver,
                layout.network.headerLength,
                layout.network.tos,
                layout.network.totalLength,
                layout.network.identification,
                layout.network.flags,
                layout.network.fragOffset,
                layout.network.timeToLive,
                layout.network.protocol,
                layout.network.srcAddress,
                layout.network.dstAddress
            },
            layout.network.hdrChecksum,
            HashAlgorithm.csum16
        );
    }
}

/*************************************************************************
*** D E P A R S E R ***
*************************************************************************/

control FrameDeparser(packet_out p, in PacketLayout layout) {
    apply {
        p.emit(layout.link);
        p.emit(layout.network);
        p.emit(layout.transport);
    }
}

/*************************************************************************
*** S W I T C H   T O P O L O G Y ***
*************************************************************************/

V1Switch(
    FrameParser(),
    VerifierBlock(),
    IngressProcessor(),
    EgressProcessor(),
    ChecksumUpdater(),
    FrameDeparser()
) main;